/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2022-2022. All rights reserved.
 * Description: 性能评估器相关源文件
 * Author: qingjicheng
 * Create: 2022-8-1
 */

#include "connected_graph_analyzer.h"
#include "evaluator_log_def.h"

namespace HcclEvaluator {
using namespace std;
ConnectedGraphAnalyzer::ConnectedGraphAnalyzer() : opFuncTable_({
    {HCCL_OP_All_TO_ALL, bind(&ConnectedGraphAnalyzer::AlgSplitAllToAllV, this, std::placeholders::_1)},
    {HCCL_OP_All_GATHER, bind(&ConnectedGraphAnalyzer::AlgSplitAllGather, this, std::placeholders::_1)},
    {HCCL_OP_All_REDUCE, bind(&ConnectedGraphAnalyzer::AlgSplitAllReduce, this, std::placeholders::_1)},
    {HCCL_OP_BROADCAST, bind(&ConnectedGraphAnalyzer::AlgSplitBroadcast, this, std::placeholders::_1)},
    {HCCL_OP_REDUCE_SCATTER, bind(&ConnectedGraphAnalyzer::AlgSplitReduceScatter, this, std::placeholders::_1)},
    {HCCL_OP_REDUCE, bind(&ConnectedGraphAnalyzer::AlgSplitReduce, this, std::placeholders::_1)}
})
{
}

ConnectedGraphAnalyzer::~ConnectedGraphAnalyzer()
{
}

HcclResult ConnectedGraphAnalyzer::AnalyzeGraphTimeCost(const ConnectedGraph &graph, HcclOpInfo &op, float &cost)
{
    EVALUATOR_CHK_RET(CheckOpInfo(op));

    AlgSplitParams para({graph, op});
    std::vector<std::vector<SubgraphCalcInfo>> multiLevelCalcInfos = opFuncTable_[op.op](para);

    if (multiLevelCalcInfos.size() != graph.maxLevel_ + 1) {
        EVALUATOR_ERROR("calc info size[%lu] is diff with max level[%u]\n", multiLevelCalcInfos.size(),
            graph.maxLevel_);
        return HCCL_E_PARA;
    }

    float totalCost = 0;
    for (u32 level = 0; level <= graph.maxLevel_; level++) {
        float t = CalcLevelCost(multiLevelCalcInfos[level], graph.multiLayerSubgraphs_[level]);
        if (t < 0) {
            EVALUATOR_ERROR("calc level[%u] time cost failed\n", level);
            return HCCL_E_PARA;
        }
        totalCost += t;
    }

    cost = totalCost;
    return HCCL_SUCCESS;
}

HcclResult ConnectedGraphAnalyzer::CheckOpInfo(HcclOpInfo &op) const
{
    if (opFuncTable_.find(op.op) == opFuncTable_.end()) {
        EVALUATOR_ERROR("op value[%d] is invalid", op.op);
        return HCCL_E_PARA;
    }

    if (op.dataType < HCCL_DATA_TYPE_INT8 || op.dataType >= HCCL_DATA_TYPE_RESERVED) {
        EVALUATOR_ERROR("data type value[%d] is invalid", op.dataType);
        return HCCL_E_PARA;
    }
    return HCCL_SUCCESS;
}

float ConnectedGraphAnalyzer::CalcLevelCost(std::vector<SubgraphCalcInfo> &calcInfos,
    const std::vector<std::unique_ptr<ConnectedSubgraph>> &subgraphs)
{
    std::vector<float> timeCosts;

    for (auto &subgraph : subgraphs) {
        float t = CalcSubgraphCost(calcInfos, subgraph);
        if (t < 0) {
            return t;
        }
        timeCosts.push_back(t);
    }

    return GetArrMaxValue(timeCosts);
}

float ConnectedGraphAnalyzer::CalcSubgraphCost(std::vector<SubgraphCalcInfo> &calcInfos,
    const std::unique_ptr<ConnectedSubgraph>& subgraph)
{
    float time = 0;
    for (auto &info : calcInfos) {
        float t = 0;
        HcclResult ret = subgraph->CalcTimeCost(info, t);
        if (ret != HCCL_SUCCESS) {
            return -1.0f;
        }
        if (t < 0) {
            return t;
        }
        time += t;
    }

    return time;
}

std::vector<std::vector<SubgraphCalcInfo>> ConnectedGraphAnalyzer::AlgSplitAllToAllV(AlgSplitParams para)
{
    std::vector<std::vector<SubgraphCalcInfo>> outputVec;
    SubgraphCalcInfo calcInfo;
    calcInfo.algType = AlgType::PAIR_WISE;
    calcInfo.subgraphOpType = SubgraphOpType::ALLTOALL;
    calcInfo.count = para.opInfo.count;
    calcInfo.dtype = para.opInfo.dataType;
    std::vector<SubgraphCalcInfo> calcInfoVec = {calcInfo};
    for (u32 level = 0; level <= para.graph.maxLevel_; level++) {
        outputVec.push_back(calcInfoVec);
    }

    return outputVec;
}

std::vector<std::vector<SubgraphCalcInfo>> ConnectedGraphAnalyzer::AlgSplitAllGather(AlgSplitParams para)
{
    std::vector<std::vector<SubgraphCalcInfo>> outputVec;
    SubgraphCalcInfo calcInfo;
    calcInfo.algType = AlgType::NOT_SPECIFIED;
    calcInfo.subgraphOpType = SubgraphOpType::ALLGATHER;
    calcInfo.count = para.opInfo.count;
    calcInfo.dtype = para.opInfo.dataType;
    std::vector<SubgraphCalcInfo> calcInfoVec = {calcInfo};
    for (u32 level = 0; level <= para.graph.maxLevel_; level++) {
        outputVec.push_back(calcInfoVec);
    }

    return outputVec;
}

std::vector<std::vector<SubgraphCalcInfo>> ConnectedGraphAnalyzer::AlgSplitAllReduce(AlgSplitParams para)
{
    std::vector<std::vector<SubgraphCalcInfo>> outputVec;
    SubgraphCalcInfo calcInfo0, calcInfo1, calcInfo2;
    calcInfo0.algType = AlgType::NOT_SPECIFIED;
    calcInfo0.subgraphOpType = SubgraphOpType::REDUCE_SCATTER;
    calcInfo0.count = para.opInfo.count;
    calcInfo0.dtype = para.opInfo.dataType;

    calcInfo1.algType = AlgType::NOT_SPECIFIED;
    calcInfo1.subgraphOpType = SubgraphOpType::ALLGATHER;
    calcInfo1.count = para.opInfo.count;
    calcInfo1.dtype = para.opInfo.dataType;

    calcInfo2.algType = AlgType::NOT_SPECIFIED;
    calcInfo2.subgraphOpType = SubgraphOpType::ALLREDUCE;
    calcInfo2.count = para.opInfo.count;
    calcInfo2.dtype = para.opInfo.dataType;

    std::vector<SubgraphCalcInfo> calcInfoVec0 = {calcInfo0, calcInfo1};
    std::vector<SubgraphCalcInfo> calcInfoVec1 = {calcInfo2};
    for (u32 level = 0; level <= para.graph.maxLevel_; level++) {
        if (level == para.graph.maxLevel_) {
            outputVec.push_back(calcInfoVec1);
        } else {
            calcInfoVec0[0].count = calcInfoVec0[0].count / para.graph.subgraphNodeSize_[level];
            calcInfoVec0[1].count = calcInfoVec0[1].count / para.graph.subgraphNodeSize_[level];
            calcInfoVec1[0].count = calcInfoVec1[0].count / para.graph.subgraphNodeSize_[level];
            outputVec.push_back(calcInfoVec0);
        }
    }

    return outputVec;
}

std::vector<std::vector<SubgraphCalcInfo>> ConnectedGraphAnalyzer::AlgSplitBroadcast(AlgSplitParams para)
{
    std::vector<std::vector<SubgraphCalcInfo>> outputVec;
    SubgraphCalcInfo calcInfo0, calcInfo1, calcInfo2;
    calcInfo0.algType = AlgType::NOT_SPECIFIED;
    calcInfo0.subgraphOpType = SubgraphOpType::SCATTER;
    calcInfo0.count = para.opInfo.count;
    calcInfo0.dtype = para.opInfo.dataType;

    calcInfo1.algType = AlgType::NOT_SPECIFIED;
    calcInfo1.subgraphOpType = SubgraphOpType::ALLGATHER;
    calcInfo1.count = para.opInfo.count;
    calcInfo1.dtype = para.opInfo.dataType;

    calcInfo2.algType = AlgType::NOT_SPECIFIED;
    calcInfo2.subgraphOpType = SubgraphOpType::BROADCAST;
    calcInfo2.count = para.opInfo.count;
    calcInfo2.dtype = para.opInfo.dataType;

    std::vector<SubgraphCalcInfo> calcInfoVec0 = {calcInfo0, calcInfo1};
    std::vector<SubgraphCalcInfo> calcInfoVec1 = {calcInfo2};
    for (u32 level = 0; level <= para.graph.maxLevel_; level++) {
        if (level == para.graph.maxLevel_) {
            outputVec.push_back(calcInfoVec1);
        } else {
            calcInfoVec0[0].count = calcInfoVec0[0].count / para.graph.subgraphNodeSize_[level];
            calcInfoVec0[1].count = calcInfoVec0[1].count / para.graph.subgraphNodeSize_[level];
            calcInfoVec1[0].count = calcInfoVec1[0].count / para.graph.subgraphNodeSize_[level];
            outputVec.push_back(calcInfoVec0);
        }
    }

    return outputVec;
}

std::vector<std::vector<SubgraphCalcInfo>> ConnectedGraphAnalyzer::AlgSplitReduceScatter(AlgSplitParams para)
{
    std::vector<std::vector<SubgraphCalcInfo>> outputVec;
    SubgraphCalcInfo calcInfo;
    calcInfo.algType = AlgType::NOT_SPECIFIED;
    calcInfo.subgraphOpType = SubgraphOpType::REDUCE_SCATTER;
    calcInfo.count = para.opInfo.count;
    calcInfo.dtype = para.opInfo.dataType;
    std::vector<SubgraphCalcInfo> calcInfoVec = {calcInfo};
    for (u32 level = 0; level <= para.graph.maxLevel_; level++) {
        outputVec.push_back(calcInfoVec);
    }
    return outputVec;
}

std::vector<std::vector<SubgraphCalcInfo>> ConnectedGraphAnalyzer::AlgSplitReduce(AlgSplitParams para)
{
    std::vector<std::vector<SubgraphCalcInfo>> outputVec;
    SubgraphCalcInfo calcInfo0, calcInfo1, calcInfo2;
    calcInfo0.algType = AlgType::NOT_SPECIFIED;
    calcInfo0.subgraphOpType = SubgraphOpType::REDUCE_SCATTER;
    calcInfo0.count = para.opInfo.count;
    calcInfo0.dtype = para.opInfo.dataType;

    calcInfo1.algType = AlgType::NOT_SPECIFIED;
    calcInfo1.subgraphOpType = SubgraphOpType::GATHER;
    calcInfo1.count = para.opInfo.count;
    calcInfo1.dtype = para.opInfo.dataType;

    calcInfo2.algType = AlgType::NOT_SPECIFIED;
    calcInfo2.subgraphOpType = SubgraphOpType::REDUCE;
    calcInfo2.count = para.opInfo.count;
    calcInfo2.dtype = para.opInfo.dataType;

    std::vector<SubgraphCalcInfo> calcInfoVec0 = {calcInfo0, calcInfo1};
    std::vector<SubgraphCalcInfo> calcInfoVec1 = {calcInfo2};
    for (u32 level = 0; level <= para.graph.maxLevel_; level++) {
        if (level == para.graph.maxLevel_) {
            outputVec.push_back(calcInfoVec1);
        } else {
            calcInfoVec0[0].count = calcInfoVec0[0].count / para.graph.subgraphNodeSize_[level];
            calcInfoVec0[1].count = calcInfoVec0[1].count / para.graph.subgraphNodeSize_[level];
            calcInfoVec1[0].count = calcInfoVec1[0].count / para.graph.subgraphNodeSize_[level];
            outputVec.push_back(calcInfoVec0);
        }
    }
    return outputVec;
}
} // namespace HcclEvaluator
